import PIL
from PIL import Image
import torch
import torchvision.transforms as transforms
from model import UnetPlusPlus
import numpy as np
import torch.nn.functional as F
import cv2  # 用于高斯模糊
from collections import OrderedDict

# 加载模型
model = UnetPlusPlus(2)
# 加载权重文件
state_dict = torch.load(r"D:\PyCharm\Unet++\Unet++_flood_model_50_b16.pth", map_location=torch.device('cuda'))

# 移除 'module.' 前缀
new_state_dict = OrderedDict()
for k, v in state_dict.items():
    if k.startswith("module."):
        new_state_dict[k[7:]] = v  # 去掉 'module.' 前缀
    else:
        new_state_dict[k] = v

# 加载到模型
model.load_state_dict(new_state_dict)

model.eval()

# 定义图像变换
trans = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor()
])

# 读取图像并应用变换
img_path = (r'D:\PyCharm\Unet++\flood\images\1.jpg')
img = PIL.Image.open(img_path)

# 使用变换调整图像尺寸并转换为tensor
img_transformed = trans(img)
# 添加批次维度
img_transformed = img_transformed.unsqueeze(0)

# 模型推理
with torch.no_grad():
    output = model(img_transformed)

# 应用Sigmoid激活函数（如果模型没有自动应用Sigmoid）
output = F.sigmoid(output)  # 使输出在[0,1]范围内

# 选择输出的通道（通常选择目标类通道）
output = output[:, 0, :, :]  # 选择第一个通道

# 处理模型输出（阈值化处理）
output = (output > 0.73).float()  # 应用阈值化，将输出转为二值图像
output=1-output
# 将输出转为可显示的图像
tensor = output.cpu()
tensor = tensor.squeeze(0)  # 去除批次维度

# 转换张量为NumPy数组，并调整像素值范围
image = tensor.numpy()
image = (image * 255).astype(np.uint8)  # 将值范围从[0, 1]调整为[0, 255]

# 应用高斯模糊来平滑图像
blurred_image = cv2.GaussianBlur(image, (5, 5), 0)  # (5, 5) 为高斯核大小，0为自动计算标准差

# 创建PIL图像并显示
segmented_img = Image.fromarray(blurred_image)
segmented_img.show()  # 显示平滑后的分割图

# 显示原始图像
img.show()  # 显示原始图像
